#!/usr/bin/env python3

from flask import Flask, render_template, request, Response, abort
from flask_socketio import SocketIO
import mss
import mss.tools
from PIL import Image
import io
import threading
import time
import json
import os

# 控制库
try:
    from pynput.mouse import Controller as MouseController, Button
    from pynput.keyboard import Controller as KeyboardController
except ImportError:
    print("请安装 pynput: pip install pynput")
    exit(1)

# 读取配置文件
def load_config():
    config = {}
    if os.path.exists('.config'):
        with open('.config', 'r') as f:
            for line in f:
                line = line.strip()
                if '=' in line and not line.startswith('#'):
                    k, v = line.split('=', 1)
                    config[k.strip()] = v.strip()
    return config.get('USERNAME'), config.get('PASSWORD')

# 加载用户名密码
USERNAME, PASSWORD = load_config()

if not USERNAME or not PASSWORD:
    print("❌ 错误：.config 文件中缺少 USERNAME 或 PASSWORD")
    exit(1)

print(f"✅ 配置加载成功: 用户名='{USERNAME}', 密码='{'*' * len(PASSWORD)}'")

app = Flask(__name__)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading')

mouse = MouseController()
keyboard = KeyboardController()

clients = 0
frame_thread = None

# 检查认证（从 URL 参数提取）
def authenticate():
    path = request.path
    if f"Username={USERNAME}" in path and f"Password={PASSWORD}" in path:
        return True
    return False


# 路由：访问入口
@app.route('/api/v1/Username=<username>/Password=<password>')
def index(username, password):
    if username == USERNAME and password == PASSWORD:
        return render_template('index.html')
    else:
        abort(403)  # 禁止访问

# WebSocket 连接事件
@socketio.on('connect')
def handle_connect():
    if not authenticate():
        return False  # 拒绝连接

    global clients
    clients += 1
    print("Connented")
    socketio.emit('status', {'msg': '已连接到远程桌面'})


@socketio.on('disconnect')
def handle_disconnect():
    global clients
    clients -= 1
    print("Clients down!")


# 接收鼠标事件
@socketio.on('mouse')
def handle_mouse(data):
    if not authenticate():
        return

    action = data.get('action')
    x = float(data.get('x'))  # 归一化坐标 [0~1]
    y = float(data.get('y'))

    screen_w, screen_h = 1920, 1080  # 默认值（实际动态获取）

    with mss.mss() as sct:
        monitor = sct.monitors[1]
        screen_w = monitor["width"]
        screen_h = monitor["height"]

    real_x = int(x * screen_w)
    real_y = int(y * screen_h)

    mouse.position = (real_x, real_y)

    if action == 'click':
        button = Button.left if data.get('button') == 'left' else Button.right
        mouse.click(button)
    elif action == 'down':
        mouse.press(Button.left)
    elif action == 'up':
        mouse.release(Button.left)


# 接收键盘事件
@socketio.on('key')
def handle_key(data):
    if not authenticate():
        return

    key = data.get('key')
    pressed = data.get('pressed', True)

    try:
        if pressed:
            keyboard.press(key)
        else:
            keyboard.release(key)
    except:
        pass  # 忽略无法识别的键


# 发送帧线程
def send_frames():
    global clients
    while True:
        if clients <= 0:
            time.sleep(0.1)
            continue

        with mss.mss() as sct:
            monitor = sct.monitors[1]
            img = sct.grab(monitor)
            img_pil = Image.frombytes("RGB", img.size, img.bgra, "raw", "BGRX")

            buf = io.BytesIO()
            img_pil.save(buf, format="JPEG", quality=50)
            frame_data = buf.getvalue()

        socketio.emit('frame', frame_data)
        time.sleep(0.05)  # ~20 FPS


if __name__ == '__main__':
    # 启动帧发送线程
    frame_thread = threading.Thread(target=send_frames, daemon=True)
    frame_thread.start()

    print("Run up!")
    socketio.run(app, host='0.0.0.0', port=5000, debug=False)